import os
import shutil
import random 
random.seed(0)
 
segfilepath = r'./VOCdevkit/VOC2007/SegmentationClass'
imgpath = r'./VOCdevkit/VOC2007/JPEGImages'
trash = r'./VOCdevkit/VOC2007/TRASH'
saveBasePath = r"./VOCdevkit/VOC2007/ImageSets/Segmentation/"

#----------------------------------------------------------------------#
#   想要增加测试集修改trainval_percent
#   修改train_percent用于改变验证集的比例
#----------------------------------------------------------------------#
trainval_percent = 1
train_percent = 0.9

temp_seg = os.listdir(segfilepath)
for img in os.listdir(imgpath):
    if not img.replace('jpg', 'png') in temp_seg:
        shutil.move(os.path.join(imgpath, img), os.path.join(trash, img))
total_seg = []
for seg in temp_seg:
    if seg.endswith(".png"):
        total_seg.append(seg)

num = len(total_seg)  
list = range(num)  
tv = int(num * trainval_percent)  
tr = int(tv * train_percent)
trainval = random.sample(list,tv)  
train = random.sample(trainval,tr)  

print("train and val size",tv)
print("traub suze",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
 
for i  in list:  
    name = total_seg[i][:-4] + '\n'  
    if i in trainval:  
        ftrainval.write(name)
        if i in train:  
            ftrain.write(name)
        else:  
            fval.write(name)
    else:  
        ftest.write(name)
  
ftrainval.close()  
ftrain.close()  
fval.close()  
ftest.close()
